package com.cyh.server.service;

import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.creator.DefaultDataSourceCreator;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.cyh.server.mapper.DBQueryMapper;
import com.cyh.server.utils.DatasourceContextUtils;
import com.cyh.server.utils.SpringUtils;
import jakarta.annotation.Resource;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.stereotype.Service;

import javax.sql.DataSource;
import java.sql.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author chenyuehua
 * @description
 * @date 2025/5/28
 **/
@Service
public class DBQueryService {

    @Resource
    private DataSource dataSource;
    @Resource
    private DefaultDataSourceCreator dataSourceCreator;
    @Resource
    private DatasourceService datasourceService;

    private boolean invalidCurrentDatabase() {
        // 获取当前数据库源
        DynamicRoutingDataSource ds = (DynamicRoutingDataSource) dataSource;
        return ds.getDataSources().isEmpty();
    }

    @Tool(name = "listTables", description = "列出当前数据库源的所有表，也是测试数据源是否成功创建的测试工具")
    public String listTables() {
        if (invalidCurrentDatabase()) {
            return "当前未创建数据库源，请发送数据源相关配置信息以创建数据源";
        }
        return queryMysql("show tables;");
    }

    /**
     * 查询指定表的行数。
     * <p>
     * 该函数通过调用数据库查询接口，获取指定表的行数，并返回包含行数的字符串信息。
     *
     * @param tableName 需要查询行数的表名，不能为空。
     * @return 返回一个字符串，格式为 "表 {tableName} 的行数为: {rowCount}"，其中 {tableName} 是传入的表名，{rowCount} 是查询到的行数。
     */
    @Tool(name = "getTableRowCount", description = """
            查询指定表的行数。
            该函数通过调用数据库查询接口，获取指定表的行数，并返回包含行数的字符串信息。
            返回一个字符串，格式为 "表 {tableName} 的行数为: {rowCount}"，其中 {tableName} 是传入的表名，{rowCount} 是查询到的行数。
            """)
    public String getTableRowCount(@ToolParam(description = "需要查询行数的表名。") String tableName) {
        if (invalidCurrentDatabase()) {
            return "当前未创建数据库源，请发送数据源相关配置信息以创建数据源";
        }

        String dbType = "mysql";
        // 通过 SpringUtils 获取 DBQueryMapper 的实例，并调用其 getTableRowCount 方法查询表的行数
        int rowCount = SpringUtils.getBean(DBQueryMapper.class).getTableRowCount(dbType, tableName);

        // 返回包含表名和行数的字符串
        return "表 " + tableName + " 的行数为: " + rowCount;
    }

    @Tool(name = "querySQL", description = """
        执行sql查询。
        该函数通过嗲用数据库查询接口，调用SQL查询，并返回SQL执行结果。
        """)
    public String queryMysql(@ToolParam(description = "标准查询sql") String sql) {
        if (invalidCurrentDatabase()) {
            return "当前未创建数据库源，请发送数据源相关配置信息以创建数据源";
        }

        PreparedStatement ps;
        List<Map<String, Object>> result;
        try {
            Connection connection = dataSource.getConnection();
            ps = connection.prepareStatement(sql);
            ResultSet resultSet = ps.executeQuery();
            result = parseResultSet(resultSet);
            // System.out.println("safety result:");
            // result.forEach(System.out::println);

            return String.valueOf(result);
        } catch (SQLException ignored) {
            return "查询失败";
        }
    }



    // 动态解析ResultSet为List<Map>
    private static List<Map<String, Object>> parseResultSet(ResultSet rs) throws SQLException {
        List<Map<String, Object>> resultList = new ArrayList<>();
        ResultSetMetaData metaData = rs.getMetaData();
        int columnCount = metaData.getColumnCount();

        while (rs.next()) {
            Map<String, Object> row = new HashMap<>();
            for (int i = 1; i <= columnCount; i++) {
                String columnName = metaData.getColumnLabel(i);
                row.put(columnName, rs.getObject(i));
            }
            resultList.add(row);
        }
        return resultList;
    }

}
